# -*- coding: utf-8 -*-
# 一致性哈希，实现一个哈希环
# 传统的hash方案，如取模，会导致数据同节点的对应关系产生巨大变化。一致性哈希的作用在于尽量减小这种变化。
from typing import Any, Dict, List, Union
import mmh3
import bisect

def mmh3_hash(input:str) -> int:
    return mmh3.hash(input, signed=False)

class VNode:
    key:int
    real_node:str

    def __init__(self, key:int, real_node:str):
        self.key = key
        self.real_node = real_node
    
    def __eq__(self, other:Any):
        return self.key == other.key

    def __lt__(self, other:Any):
        return self.key < other.key
    
    def __repr__(self) -> str:
        return f'[key:{self.key}, real_node:{self.real_node}]'

    def __str__(self) -> str:
        return f'[key:{self.key}, real_node:{self.real_node}]'

class ServerNode:
    def __init__(self, name:str, ref:Any, replica:int):
        self.name = name
        self.ref = ref
        self.replica = replica
        self.vnodes = []

    def add_vnodes(self, vnode:VNode):
        self.vnodes.append(vnode)

class HashRing:

    def __init__(self, nodes:List[Union[str,ServerNode]] = [], default_replica:int=4):
        self.nodes:Dict[str,ServerNode] = {}
        self.vnodes:List[VNode] = []
        self.vnode_key_set = set()
        self.default_replica = default_replica
        self.add_nodes(nodes)

    def add_node(self, server_node:Union[str,ServerNode]):
        if isinstance(server_node, str):
            server_node = ServerNode(server_node, None, self.default_replica)
        if server_node.name in self.nodes:
            raise Exception(f'server node {server_node.name} already in hashring')
        self.nodes[server_node.name] = server_node
        self.__add_virtual_nodes(server_node)
        self.__sort()

    def add_nodes(self, server_nodes:List[Union[str,ServerNode]]):
        for server_node in server_nodes:
            if isinstance(server_node, str):
                server_node = ServerNode(server_node, None, self.default_replica)
            if server_node.name in self.nodes:
                raise Exception(f'server node {server_node.name} already in hash ring')
            self.nodes[server_node.name] = server_node
            self.__add_virtual_nodes(server_node)
        self.__sort()

    def remove_node(self, name:str):
        if name not in self.nodes:
            raise Exception(f'server node {name} not in hash ring')
        self.__remove_virtual_nodes(self.nodes[name])
        self.__sort()
        del self.nodes[name]

    def hash(self, input:Union[str,int]) -> ServerNode:
        if len(self.vnodes) < 1:
            raise Exception('no vnodes in hash ring')
        if isinstance(input, str):
            input = mmh3_hash(input)
        target = self.find_vnode(input, self.vnodes)
        return self.nodes[target.real_node]

    def __add_virtual_nodes(self, real_node:ServerNode):
        for i in range(real_node.replica):
            to_be_hash = f'{real_node.name}#{i}'
            v_key = mmh3_hash(to_be_hash)
            #处理可能的冲突
            while v_key in self.vnode_key_set:
                to_be_hash += '#1'
                v_key =  mmh3_hash(to_be_hash)
            self.vnode_key_set.add(v_key)
            vnode = VNode(v_key, real_node.name)
            self.vnodes.append(vnode)
            real_node.add_vnodes(vnode)
    
    def __remove_virtual_nodes(self,real_node:ServerNode):
        for v_node in real_node.vnodes:
            idx = bisect.bisect(self.vnodes,v_node)
            self.vnodes.pop(idx - 1)
            self.vnode_key_set.remove(v_node.key)

    def __sort(self):
        self.vnodes.sort(key = lambda x:x.key)

    #二分查找的变体，寻找一个key归属的虚拟节点
    @classmethod
    def find_vnode(cls, key:int, vnodes:List[VNode]) -> VNode:
        if key <= vnodes[0].key or key > vnodes[-1].key:
            return vnodes[0]
        left = 0
        right = len(vnodes) - 1
        while left + 1 < right:
            mid = int((left+right)/2)
            if vnodes[mid].key == key:
                return vnodes[mid]
            elif vnodes[mid].key < key:
                left = mid
            else:
                right = mid
        return vnodes[right]
        
if __name__ == '__main__':
    hr = HashRing(["node1",'node2','node3'], 3)
    print(hr.vnodes)
    assert hr.find_vnode(1, hr.vnodes) == hr.vnodes[0]
    assert hr.find_vnode(435927186, hr.vnodes) == hr.vnodes[0]
    assert hr.find_vnode(3901558265, hr.vnodes) == hr.vnodes[0]

    assert hr.find_vnode(606617729, hr.vnodes) == hr.vnodes[1]

    assert hr.find_vnode(606617730, hr.vnodes) == hr.vnodes[1]

    assert hr.find_vnode(606617732, hr.vnodes) == hr.vnodes[2]

    assert hr.find_vnode(607538691, hr.vnodes) == hr.vnodes[2]

    assert hr.find_vnode(607538692, hr.vnodes) == hr.vnodes[3]

    assert hr.find_vnode(1070073501, hr.vnodes) == hr.vnodes[3]

    assert hr.find_vnode(3901558263, hr.vnodes) == hr.vnodes[8]

    assert hr.find_vnode(3901558264, hr.vnodes) == hr.vnodes[8]
